package com.meituan.hive.udf.generic;

import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.exec.UDFArgumentLengthException;
import org.apache.hadoop.hive.ql.exec.UDFArgumentTypeException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDF;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDFUtils;
import org.apache.hadoop.hive.serde.serdeConstants;
import org.apache.hadoop.hive.serde2.objectinspector.ListObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorConverters;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * Created by zhaoxiang on 2017/8/9.
 * 合并两个数组并去重
 * func(array1, array2)
 */
public class GenericUDFArrayAddDedup extends GenericUDF {

    private transient ObjectInspectorConverters.Converter[] converters;
    private final List<Object> ret = new ArrayList<Object>();
    private final Set<Object> retSet = new HashSet<>();
    private transient ObjectInspector[] argumentOIs;

    protected String getFuncName() {
        return getClass().getSimpleName().substring(10).toLowerCase();
    }

    protected void checkArgsSize(ObjectInspector[] arguments, int min, int max)
            throws UDFArgumentLengthException {
        if (arguments.length < min || arguments.length > max) {
            StringBuilder sb = new StringBuilder();
            sb.append(getFuncName());
            sb.append(" requires ");
            if (min == max) {
                sb.append(min);
            } else {
                sb.append(min).append("..").append(max);
            }
            sb.append(" argument(s), got ");
            sb.append(arguments.length);
            throw new UDFArgumentLengthException(sb.toString());
        }
    }

    @Override
    public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumentException {
        GenericUDFUtils.ReturnObjectInspectorResolver returnOIResolver;
        returnOIResolver = new GenericUDFUtils.ReturnObjectInspectorResolver(true);

        checkArgsSize(arguments, 2, Integer.MAX_VALUE);

        switch(arguments[0].getCategory()) {
            case LIST:
                if(!((ListObjectInspector)(arguments[0])).getListElementObjectInspector()
                        .getCategory().equals(ObjectInspector.Category.UNION)) {
                    break;
                }
            default:
                throw new UDFArgumentTypeException(0, "Argument 1"
                        + " of function ArrayAddDedup must be " + serdeConstants.LIST_TYPE_NAME
                        + ", and element type should be either primitive, list, struct, or map, " +
                        "but " + arguments[0].getTypeName() + " was found.");
        }

        ObjectInspector elementObjectInspector =
                ((ListObjectInspector)(arguments[0])).getListElementObjectInspector();
        argumentOIs = arguments;
        converters = new ObjectInspectorConverters.Converter[arguments.length];
        ObjectInspector returnOI = returnOIResolver.get(elementObjectInspector);
        converters[0] = ObjectInspectorConverters.getConverter(elementObjectInspector, returnOI);

        return ObjectInspectorFactory.getStandardListObjectInspector(returnOI);
    }

    @Override
    public Object evaluate(GenericUDF.DeferredObject[] arguments) throws HiveException {
        retSet.clear();
        ret.clear();

        Object array1 = arguments[0].get();
        Object array2 = arguments[1].get();
        ListObjectInspector arrayOI = (ListObjectInspector) argumentOIs[0];
        List retArray1 = (List) arrayOI.getList(array1);
        List retArray2 = (List) arrayOI.getList(array2);

        if (array1 == null && array2 == null) {
            return ret;
        } else if (array1 == null) {
            return retArray2;
        }else if(array2 == null){
            return retArray1;
        }


        retSet.addAll(retArray1);
        retSet.addAll(retArray2);

        ret.addAll(retSet);

        //for (int i = 0; i < retArray.size(); i++) {
        //    ret.add(converters[0].convert(retArray.get(i)));
        //}
        return ret;
    }

    @Override
    public String getDisplayString(String[] children) {
        assert (children.length == 1);
        return getStandardDisplayString("array_add", children);
    }
}
